import os
import math

import bpy
from bpy.props import (
    BoolProperty,
    FloatProperty,
    StringProperty,
    EnumProperty,
)
from bpy_extras.io_utils import (
    ImportHelper,
    ExportHelper,
    orientation_helper,
    path_reference_mode,
    axis_conversion,
)


from mathutils import Matrix, Vector, Color
from bpy_extras import io_utils, node_shader_utils

from bpy_extras.wm_utils.progress_report import (
	ProgressReport,
	ProgressReportSubstep,
)



def mesh_triangulate(me):
	import bmesh
	bm = bmesh.new()
	bm.from_mesh(me)
	bmesh.ops.triangulate(bm, faces=bm.faces)
	bm.to_mesh(me)
	bm.free()

def pos_object_location_str(ob):
	pos = ' pos = \"' 
	pos += ' ' + str(ob.location.x) 
	pos += ' ' + str(ob.location.y)
	pos += ' ' + str(ob.location.z)
	pos += '\" '
	return pos

def name_object_name_str(ob, name='name', 后缀=''):
	return ' ' + name + ' =\"' + str(ob.name) + 后缀 + '\" '



def ob_write_mesh(objects, path, fw):
	bpy.ops.export_mesh.stl(filepath=path, batch_mode='OBJECT')

	fileName = os.path.basename(path)
	fw("<asset> \n")

	for i, ob in enumerate(objects):
		if ob.type == 'MESH' :
			name = fileName + ob.name
			name = bpy.path.clean_name(name)
			fw( '	<mesh file = \'' + name + '.stl\' '  + 'name = \'' + ob.name + '\' /> \n ' )
	fw("</asset> \n")
	
def ob_write_刚体链接(rigid_body_constraint, ob):
	pos = pos_object_location_str(ob.rigid_body_constraint.object1)

	joint = ''
	限制值 = ''
	if rigid_body_constraint.use_breaking :
		限制值 = str(rigid_body_constraint.breaking_threshold)
		限制值 = ' limited=\"true\" solimplimit=\"' + 限制值 + ' ' + 限制值 + ' ' + 限制值 + '\"' + ' damping = \"' + 限制值 + '\"'
		
	if rigid_body_constraint.type == 'POINT' :
		joint = '	<joint type = \"ball\" ' + pos + 限制值 + ' /> \n'
	elif rigid_body_constraint.type == 'HINGE' :
		joint = '	<joint type = \"hinge\" ' + pos + 限制值 + ' /> \n'
	elif rigid_body_constraint.type == 'PISTON' :
		joint = '	<joint type = \"slide\" ' + pos + 限制值 + ' /> \n'
	elif rigid_body_constraint.type == 'FIXED' :
		joint = '	<joint type = \"free\" ' + pos + 限制值 + ' /> \n'
	elif rigid_body_constraint.type == 'GENERIC' or rigid_body_constraint.type == 'GENERIC_SPRING' :
		轴限定 = ''
		if rigid_body_constraint.use_limit_ang_x :
			轴限定 = name_object_name_str(ob, name='name', 后缀='_hip_x') #'name = \"' + ob.name + '_hip_x\" '
			轴限定 += 'axis=\"1 0 0\" ' + ' range=\"' + str(math.degrees(rigid_body_constraint.limit_ang_x_lower)) + ' ' + str(math.degrees(rigid_body_constraint.limit_ang_x_upper)) +'\" '
			joint += '	<joint type = \"hinge\" ' + 轴限定 + pos + 限制值 + ' /> \n'
		if rigid_body_constraint.use_limit_ang_y :
			轴限定 = name_object_name_str(ob, name='name', 后缀='_hip_y')
			轴限定 += 'axis=\"0 1 0\" ' + ' range=\"' + str(math.degrees(rigid_body_constraint.limit_ang_y_lower)) + ' ' + str(math.degrees(rigid_body_constraint.limit_ang_y_upper)) +'\" '
			joint += '	<joint type = \"hinge\" ' + 轴限定 + pos + 限制值 + ' /> \n'
		if rigid_body_constraint.use_limit_ang_z :
			轴限定 = name_object_name_str(ob, name='name', 后缀='_hip_z')
			轴限定 += 'axis=\"0 0 1\" ' + ' range=\"' + str(math.degrees(rigid_body_constraint.limit_ang_z_lower)) + ' ' + str(math.degrees(rigid_body_constraint.limit_ang_z_upper)) +'\" '
			joint += '	<joint type = \"hinge\" ' + 轴限定 + pos + 限制值 + ' /> \n'

	return joint


def ob_write(ob):
	xml = ''
	gepm_type = "mesh"
	位置 = ' pos = \"' + str(ob.location.x) + ' ' + str(ob.location.y) + ' ' + str(ob.location.z) + '\"'
	旋转 = ' euler = \"' + str(ob.rotation_euler[0]) + ' ' + str(ob.rotation_euler[1]) + ' ' + str(ob.rotation_euler[2]) + '\"'
	大小 = ' size = \"' + str(ob.scale.x) + ' ' + str(ob.scale.y) + ' ' + str(ob.scale.z) + '\"'

	链接XML = ''
	if ob.rigid_body_constraint != None :
		#pos = '' 
		if ob.rigid_body_constraint.object1 != None :
			#pos = pos_object_location_str(ob.rigid_body_constraint.object1)
			链接XML += ob_write_刚体链接(ob.rigid_body_constraint, ob)
		#limited = "true" solimplimit="0.1,0.1,0.1" 
		#链接XML += joint
		位置 = ' pos = \" 0 0 0 \"'

	刚体参数 = ' condim=\"1\" friction=\"0.7\" solimp=\".9 .99 .003\" solref=\"0.015 1\" '
	xml = "	<body " + name_object_name_str(ob) + 位置 + 旋转 + "> \n";
	if ob.rigid_body != None :
		if ob.rigid_body.type == 'ACTIVE' :
			xml += "	<freejoint /> \n";
	else:
		pass
	xml += 链接XML
	
	
	#重量
	xml += "	<geom type = \"" + gepm_type + "\" " + 'mesh=\'' + ob.name + '\' ' + 大小 + "/> \n";

	#肌腱描点
	for mod in ob.modifiers:
		if mod.type == 'HOOK' and mod.object != None:
			pos = pos_object_location_str(mod.object)
			name = name_object_name_str(mod.object)
			xml += '		<site ' + name + pos + '  rgba=\"0 .7 0 1\" size=\"0.01\"/> \n'
			

	for ob in ob.children:
		if ob.type == 'MESH':
			xml += ob_write(ob)[0]

	xml += "	</body> \n";

	资源 = ''
	return xml, 资源


def ob_write_肌腱():
	xml = ''
	collection = bpy.data.collections.get('肌腱')
	if collection != None :
		xml += '<tendon> \n'
		for ob in collection.objects:
			if ob.type == 'CURVE' :
				xml += '<spatial '+ name_object_name_str(ob) +' width=\"0.005\"> \n'
				for mod in ob.modifiers:
					if mod.type == 'HOOK' and mod.object != None:
						xml += '<site '+ name_object_name_str(mod.object, 'site') +'/> \n'
						#xml += '<geom geom=\"shoulder\"/> \n'
				#bpy.data.curves:
				xml += '</spatial> \n'

		xml += '</tendon> \n'
	return xml

class ExportMoJoCo(bpy.types.Operator, ExportHelper):
	bl_idname = "export_scene.xml"
	bl_label = '导出 MuJoCo XML'
	bl_options = {'PRESET'}

	filename_ext = ".xml"
	filepath = StringProperty(subtype='FILE_PATH')

	def execute(self, context):
		dirname = os.path.dirname(self.filepath)
		#relpath = bpy.data.filepath
		#if context.preferences.filepaths.use_relative_paths :
		#	relpath = os.path.dirname(bpy.data.filepath)
		path = os.path.join(dirname, self.filepath)

		with open(path, "w", encoding="utf8", newline="\n") as f:
			#盆骨 = context.selected_objects
			盆骨 = context.object
			#objects = bpy.data.objects

			fw = f.write
			fw("<mujoco> \n")

			default = '<default> \n'
			default += '<muscle ctrllimited=\"true\" ctrlrange=\"0 1\"/> \n'
			default += '</default> \n'
			fw(default)

			模型资源保存路径 = os.path.splitext(path)[0]
			ob_write_mesh(bpy.data.objects, 模型资源保存路径, fw)

			fw("	<worldbody> \n");
			fw('	<geom name=\"ground\" type=\"plane\" size=\"0 0 1\" pos=\"0 0 0\" quat=\"1 0 0 0\" condim=\"1\"/> \n')
			if 盆骨.type == 'MESH':
				xml = ob_write(盆骨)
				fw(xml[0])

			fw("	</worldbody> \n");

			fw(ob_write_肌腱())

			fw("</mujoco> \n")

		print("dirname:", path)

		return {'FINISHED'}


classes = (
	ExportMoJoCo,
)



def menu_func_export(self, context):
	self.layout.operator(ExportMoJoCo.bl_idname, text="MuJoCo (.xml)")

def register():
	for cls in classes:
		bpy.utils.register_class(cls)

	#bpy.types.TOPBAR_MT_file_import.append(menu_func_import)
	bpy.types.TOPBAR_MT_file_export.append(menu_func_export)


def unregister():
	#bpy.types.TOPBAR_MT_file_import.remove(menu_func_import)
	bpy.types.TOPBAR_MT_file_export.remove(menu_func_export)

	for cls in classes:
		bpy.utils.unregister_class(cls)


#if __name__ == "__main__":
#	register()
register()
